package com.smart.milvus.autoconfigure;

import com.smart.milvus.base.*;
import com.smart.milvus.bean.*;
import io.milvus.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

@Slf4j
@Configuration(proxyBeanMethods = false)
@EnableConfigurationProperties(MilvusProperties.class)
public class IMilvusAutoConfiguration implements InitializingBean {

    private final MilvusProperties properties;

    public IMilvusAutoConfiguration(MilvusProperties properties) {
        this.properties = properties;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
    }

    @Bean
    @ConditionalOnMissingBean
    public IdentifierGenerator identifierGenerator() {
        return DefaultIdentifierGenerator.getInstance();
    }


    @Bean
    @ConditionalOnMissingBean
    public AnnotationHandler annotationHandler() {
        return new AnnotationHandler() {
        };
    }

    @Bean
    @ConditionalOnMissingBean
    public DefaultSpElExecutor spElHandler() {
        return new DefaultSpElExecutor();
    }

    @Bean
    @ConditionalOnMissingBean
    public PostInitCollectionInfoHandler postInitCollectionInfoHandler() {
        return new PostInitCollectionInfoHandler() {
        };
    }

    @Bean
    @ConditionalOnMissingBean
    public UnknownTypeHandler defaultEmbeddingTypeHandler() {
        return new UnknownTypeHandler();
    }

    @Bean
    @ConditionalOnMissingBean
    public JsonTypeHandler defaultJsonTypeHandler() {
        return new JsonTypeHandler() {
        };
    }

    @Bean
    @ConditionalOnMissingBean
    public MilvusClient milvusClient() {
        if (properties == null) {
            return null;
        }
        ConnectParam.Builder builder = ConnectParam.newBuilder();
        if (properties.getHost() != null) {
            builder.withHost(properties.getHost());
        }
        if (properties.getPort() != null) {
            builder.withPort(properties.getPort());
        }
        if (properties.getUserName() != null) {
            builder.withAuthorization(properties.getUserName(), properties.getPassword());
        }
        if (properties.getConnectUri() != null) {
            builder.withUri(properties.getConnectUri());
        }
        if (properties.getConnectTimeout() != null) {
            builder.withConnectTimeout(properties.getConnectTimeout(), TimeUnit.MILLISECONDS);
        }
        if (properties.getRpcDeadline() != null) {
            builder.withRpcDeadline(properties.getRpcDeadline(), TimeUnit.MILLISECONDS);
        }
        if (properties.getDatabase() != null) {
            builder.withDatabaseName(properties.getDatabase());
        }
        if (properties.getSecure() != null) {
            builder.withSecure(properties.getSecure());
        }
        if (properties.getKeepAliveTime() != null) {
            builder.withKeepAliveTime(properties.getKeepAliveTime(), TimeUnit.MILLISECONDS);
        }
        if (properties.getIdleTimeout() != null) {
            builder.withIdleTimeout(properties.getIdleTimeout(), TimeUnit.MILLISECONDS);
        }
        if (properties.getToken() != null) {
            builder.withToken(properties.getToken());
        }

        return new MilvusServiceClient(builder.build());
    }

    @Bean
    @ConditionalOnMissingBean
    public MilvusClientService milvusClientService(MilvusClient milvusClient,
                                                   ApplicationContext applicationContext) {
        GlobalConfig globalConfig = new GlobalConfig();
        globalConfig.setMilvusProperties(properties);
        this.getBeanThen(DefaultSpElExecutor.class, globalConfig::setDefaultSpElExecutor);
        this.getBeanThen(AnnotationHandler.class, globalConfig::setAnnotationHandler);
        this.getBeanThen(IdentifierGenerator.class, globalConfig::setIdentifierGenerator);
        this.getBeanThen(PostInitCollectionInfoHandler.class, globalConfig::setPostInitCollectionInfoHandler);
        this.getBeanThen(JsonTypeHandler.class, globalConfig::setJsonTypeHandler);
        this.getBeanThen(DynamicCollectionHandler.class, globalConfig::setDynamicCollectionHandler);
        GlobalConfig.CTX = applicationContext;
        GlobalConfig.INSTANCE = globalConfig;
        return new MilvusClientService(milvusClient);
    }

    /**
     * 检查spring容器里是否有对应的bean,有则进行消费
     *
     * @param clazz    class
     * @param consumer 消费
     * @param <T>      泛型
     */
    private <T> void getBeanThen(Class<T> clazz, Consumer<T> consumer) {
        if (GlobalConfig.CTX.getBeanNamesForType(clazz, false, false).length > 0) {
            consumer.accept(GlobalConfig.CTX.getBean(clazz));
        }
    }
}
